#pragma once
#include "matrix_mem_view.hpp"
#include <thread>
#include <vector>
#include <cstring>

constexpr int THRESHOLD_SHIFT = 9;

inline void copy_2d_block(MatrixMemView destination, MatrixMemView source, 
                         size_t offset, size_t length, int32_t data_size) {
    destination.next(offset);
    source.next(offset);
    source.copyTo(destination, length, data_size);
}

template<typename DestinationType, typename SourceType>
inline void assign_2d_block(MatrixMemView destination, MatrixMemView source, 
                           size_t offset, size_t length) {
    destination.next(offset);
    source.next(offset);
    source.assignTo<DestinationType, SourceType>(destination, length);
}

void threaded_2d_copy(void* destination, void* source, 
                     int32_t row_count, int32_t column_count,
                     size_t dest_stride_0, size_t dest_stride_1,
                     size_t src_stride_0, size_t src_stride_1, 
                     int32_t data_size, int32_t thread_count);

template<typename DestinationType, typename SourceType>
inline void threaded_2d_assign(void* destination, void* source, 
                              int32_t row_count, int32_t column_count, 
                              size_t dest_stride_0, size_t dest_stride_1,
                              size_t src_stride_0, size_t src_stride_1, 
                              int32_t thread_count) {
    MatrixMemView view_destination(destination, row_count, column_count, 
                                 dest_stride_0, dest_stride_1);
    MatrixMemView view_source(source, row_count, column_count, 
                            src_stride_0, src_stride_1);
    
    size_t total_length = static_cast<size_t>(row_count) * column_count;
    size_t block_length;
    size_t offset = 0;
    int32_t remainder;
    
    if (thread_count > static_cast<int32_t>(total_length >> THRESHOLD_SHIFT)) {
        thread_count = static_cast<int32_t>(total_length >> THRESHOLD_SHIFT);
    }
    if (thread_count <= 1) {
        assign_2d_block<DestinationType, SourceType>(view_destination, view_source, 
                                                    offset, total_length);
        return;
    }
    
    block_length = total_length / thread_count;
    remainder = total_length % thread_count;
    
    std::vector<std::thread> threads;
    threads.reserve(thread_count);
    
    for (int32_t i = 0; i < thread_count - remainder; i++) {
        threads.emplace_back(assign_2d_block<DestinationType, SourceType>, 
                           view_destination, view_source, offset, block_length);
        offset += block_length;
    }
    
    block_length++;
    for (int32_t i = thread_count - remainder; i < thread_count; i++) {
        threads.emplace_back(assign_2d_block<DestinationType, SourceType>, 
                           view_destination, view_source, offset, block_length);
        offset += block_length;
    }
    
    for (auto& thread : threads) {
        thread.join();
    }
}